#!/usr/bin/python3
# -*- coding: utf-8 -*-

"""
# @version: v1.0
# @author : wlis
# @Email : 19688513@qq.com
# @Project : g-carbon-bio
# @File : trade_date_calendar.py
# @Software: PyCharm
# @time: 2025/2/4 17:32
# @description : 获取历史交易日期数据  调用接口 ak.tool_trade_date_hist_sina()
"""

import akshare as ak
from service.csv_help import CSVHelp
from datetime import datetime
import time
from pandas.tseries.offsets import CustomBusinessDay

class TradeDateCalendar:
    def __init__(self, folder_name='service/data', file_name='trade_dates.csv'):
        """
        初始化 TradeDateCalendar 类

        Args:
            folder_name (str): 存放 CSV 文件的文件夹名称，默认为 'data'
            file_name (str): 要存储的文件名，默认为 'trade_dates.csv'
        """
        self.csv_helper = CSVHelp(folder_name)
        self.file_name = file_name

    def fetch_and_store_trade_dates(self):
        """
        从 Akshare 获取交易日期数据并存储到 CSV 文件
        """
        # 获取交易日期数据
        tool_trade_date_hist_sina_df = ak.tool_trade_date_hist_sina()
        # 写入到 CSV 文件
        self.csv_helper.write_to_csv(tool_trade_date_hist_sina_df, self.file_name)

    def read_trade_dates(self):
        """
        从 CSV 文件读取交易日期数据，并返回日期列表。

        Returns:
            list: 包含交易日期的列表，格式为 ["YYYY-MM-DD", ...]
        """
        df = self.csv_helper.read_csv(self.file_name)
        if df is not None and 'trade_date' in df.columns:
            # 将 trade_date 列转换为列表
            return df['trade_date'].tolist()
        return None

    def is_tool_trading_day(self, date):
        """
        判断给定日期是否为工具交易日。
        :param date: datetime.date 对象，表示要判断的日期。
        :return: 如果是工具交易日，返回 True；否则返回 False。
        """
        if date is None:
            raise ValueError("日期不能为空")

        # 获取工具交易日期列表
        tool_days = self.read_trade_dates()

        if tool_days is not None:
            # 将工具交易日期转换为日期格式
            tool_days = [datetime.strptime(day, "%Y-%m-%d").date() for day in tool_days]
            return date in tool_days
        else:
            # 如果没有获取到任何工具交易日期，判断该日期是否为工作日
            return date.weekday() < 5  # 返回值小于 5，表示为周一到周五（工作日）

    def get_start_end_date(self, trade_date: str, y_day: int):
        """
        根据给定的交易日期和天数，计算起始日期和结束日期。

        Args:
            trade_date (str): 交易日期，格式为 "YYYY-MM-DD"
            y_day (int): 需要推算的天数

        Returns:
            tuple: 起始日期和结束日期，格式为 "YYYYMMDD"
        """
        start_time = time.time()  # 记录开始时间
        print("开始获取交易日期数据...")

        trade_dates = self.read_trade_dates()  # 返回交易日期列表

        if trade_dates is not None and trade_date in trade_dates:
            index = trade_dates.index(trade_date)
            # 计算开始日期和结束日期
            start_index = max(0, index - y_day)
            end_index = index

            start_date = trade_dates[start_index]
            end_date = trade_dates[end_index]

            # 转换成 YYYYMMDD 格式
            start_date = datetime.strptime(start_date, "%Y-%m-%d").strftime("%Y%m%d")
            end_date = datetime.strptime(end_date, "%Y-%m-%d").strftime("%Y%m%d")

            end_time = time.time()  # 记录结束时间
            print(f"获取交易日期数据完成。运行时间: {end_time - start_time:.2f}秒")
            return start_date, end_date  # 返回起始日期和结束日期
        else:
            # 将 trade_date 转换为 datetime.date 类型进行工作日推算
            trade_date_obj = datetime.strptime(trade_date, "%Y-%m-%d").date()
            cbd = CustomBusinessDay(weekmask='Mon Tue Wed Thu Fri')

            # 向前推算 y_day 个工作日
            start_date_obj = trade_date_obj - y_day * cbd
            start_date = start_date_obj.strftime("%Y%m%d")
            end_date = trade_date_obj.strftime("%Y%m%d")

            end_time = time.time()  # 记录结束时间
            print(f"获取交易日期数据完成。运行时间: {end_time - start_time:.2f}秒")
            return start_date, end_date  # 返回推算的日期

    def get_offset_trade_dates(self, trade_date_str: str, offset_days: int):
        """
        获取给定交易日期的偏移交易日期。

        Args:
            trade_date_str (str): 输入交易日期，格式为 "YYYY-MM-DD"
            offset_days (int): 偏移的天数

        Returns:
            tuple: 偏移后的日期，格式为 (previous_trade_date, next_trade_date)
        """
        if trade_date_str is None:
            raise ValueError("日期不能为空")

        # 读取交易日期列表
        tool_days = self.read_trade_dates()

        if tool_days is not None and trade_date_str in tool_days:
            # 找到交易日期的索引
            index = tool_days.index(trade_date_str)

            # 计算前后偏移的索引
            previous_index = index - offset_days
            next_index = index + offset_days

            # 获取前后交易日期，确保索引在范围内
            previous_trade_date = tool_days[previous_index] if previous_index >= 0 else None
            next_trade_date = tool_days[next_index] if next_index < len(tool_days) else None

            # 格式化为 "%Y%m%d"
            if previous_trade_date:
                previous_trade_date = datetime.strptime(previous_trade_date, "%Y-%m-%d").strftime("%Y%m%d")
            if next_trade_date:
                next_trade_date = datetime.strptime(next_trade_date, "%Y-%m-%d").strftime("%Y%m%d")

            return previous_trade_date, next_trade_date
        else:
            raise ValueError("给定的交易日期不在交易日期列表中")

# 示例用法
# if __name__ == "__main__":
#     # 初始化 TradeDateCalendar
#     trade_calendar = TradeDateCalendar()
#
#     calendar = TradeDateCalendar()
#     input_date = "2023-10-25"  # 输入的交易日期
#     offset_days = 1  # 偏移的天数
#
#     try:
#         previous_trade_date, next_trade_date = calendar.get_offset_trade_dates(input_date, offset_days)
#         print(f"输入日期: {input_date}")
#         print(f"前一个交易日期: {previous_trade_date}")
#         print(f"后一个交易日期: {next_trade_date}")
#     except ValueError as e:
#         print(e)
#
#     # 获取并存储交易日期数据
#     trade_calendar.fetch_and_store_trade_dates()
#
#     # 判断某一天是否为交易日
#     test_date = datetime.date(2023, 10, 10)
#     is_trading_day = trade_calendar.is_tool_trading_day(test_date)
#     print(f"{test_date} 是否为交易日: {is_trading_day}")
#
#     # 获取起始日期和结束日期
#     start_date, end_date = trade_calendar.get_start_end_date("1990-12-28", 5)
#     print(f"起始日期: {start_date}, 结束日期: {end_date}")